import numpy as np
from scipy.optimize import minimize_scalar

# ============ 场景常量 ============
g = 9.81  # m/s^2
R = 10.0  # 有效遮蔽半径（云团中心 10 m）
v_sink = 3  # 云团中心下沉速度 m/s
valid_win = 20  # 起爆后有效 20 s
T = np.array([0, 200, 0])  # 真目标中心
Of = np.array([0, 0, 0])  # 假目标（原点）

# 导弹 M1
M0 = np.array([20000, 0, 2000])
uM = (Of - M0) / np.linalg.norm(Of - M0)  # 朝向假目标
vM = 300 * uM  # 速度 300 m/s
L = np.linalg.norm(Of - M0)
t_hit = L / 300  # 抵达假目标时间 ~ 67 s

# 三架 UAV 初始位置（等高度直线飞行，z 保持不变）
FY = np.array([
    [17800, 0, 1800],  # FY1
    [12000, 1400, 1400],  # FY2
    [6000, -3000, 700]  # FY3
])
v_min = 70
v_max = 140

# ============ 搜索网格（可调） ============
angles = np.deg2rad(np.arange(-180, 180, 22.5))  # 航向角（相对 +x 逆时针）
speeds = np.arange(70, 150, 10)  # 速度
t_grid = np.arange(2, t_hit - 2, 2)  # 候选起爆时刻（全局时间）
tau_set = np.array([2.4, 3.0, 3.6, 4.2])  # 候选引信延时
dt_eval = 0.02  # 评估时间步长

# ============ 预计算导弹/视线量 ============
ts = np.arange(0, t_hit - 1e-6, dt_eval)  # 避免 t_hit 处除零
N = len(ts)
Mx = M0[0] + vM[0] * ts
My = M0[1] + vM[1] * ts
Mz = M0[2] + vM[2] * ts
Ux = T[0] - Mx
Uy = T[1] - My
Uz = T[2] - Mz
UU = Ux ** 2 + Uy ** 2 + Uz ** 2

# ============ 贪心三弹（跨三架 UAV） ============
dmin_all = np.full(N, np.inf)  # 当前联合最小距离
uav_used = np.zeros(3, dtype=bool)  # 每架 UAV 仅用一次


# 定义计划结构
class Plan:
    def __init__(self):
        self.uav = 0
        self.theta = 0
        self.v = 0
        self.t_rel = 0
        self.tau = 0
        self.t_det = 0
        self.D = np.zeros(3)
        self.C0 = np.zeros(3)


plan = []

for shot in range(3):
    best_gain = -np.inf
    best_info = Plan()
    best_dmin_new = None

    for iu in range(3):
        if uav_used[iu]:
            continue

        F0 = FY[iu, :]
        z0 = F0[2]

        for ia in range(len(angles)):
            theta = angles[ia]
            uxF = np.cos(theta)
            uyF = np.sin(theta)

            for iv in range(len(speeds)):
                vF = speeds[iv]

                for id in range(len(t_grid)):
                    t_det = t_grid[id]

                    for it in range(len(tau_set)):
                        tau = tau_set[it]
                        t_rel = t_det - tau

                        if t_rel < 0:
                            continue

                        # 起爆点 = UAV 位置(t_det) 与自由落体
                        x_det = F0[0] + vF * uxF * t_det
                        y_det = F0[1] + vF * uyF * t_det
                        z_det = z0 - 0.5 * g * tau ** 2

                        # 云团中心（仅在有效窗内）
                        valid = (ts >= t_det) & (ts <= t_det + valid_win)
                        Cz = z_det - v_sink * (ts - t_det)

                        # 该弹对"导弹->真目标"线段的最小距离 d_k(t)
                        Wx = x_det - Mx
                        Wy = y_det - My
                        Wz = Cz - Mz

                        WU = Wx * Ux + Wy * Uy + Wz * Uz
                        lam = WU / UU
                        lam = np.clip(lam, 0, 1)

                        Px = Mx + lam * Ux
                        Py = My + lam * Uy
                        Pz = Mz + lam * Uz

                        dx = x_det - Px
                        dy = y_det - Py
                        dz = Cz - Pz
                        d_k = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2)
                        d_k[~valid] = np.inf

                        # 与现有联合取最小
                        dmin_new = np.minimum(dmin_all, d_k)
                        gain = np.sum(dmin_new <= R) * dt_eval

                        if gain > best_gain:
                            best_gain = gain
                            best_dmin_new = dmin_new.copy()

                            best_info = Plan()
                            best_info.uav = iu
                            best_info.theta = theta
                            best_info.v = vF
                            best_info.t_rel = t_rel
                            best_info.tau = tau
                            best_info.t_det = t_det
                            best_info.D = np.array([
                                F0[0] + vF * uxF * t_rel,
                                F0[1] + vF * uyF * t_rel,
                                z0
                            ])
                            best_info.C0 = np.array([x_det, y_det, z_det])

    # 接受该次最优
    if best_gain == -np.inf:
        print(f'警告: 第 {shot + 1} 枚未找到可行解。')
        break

    dmin_all = best_dmin_new
    uav_used[best_info.uav] = True
    plan.append(best_info)

total_cov = np.sum(dmin_all <= R) * dt_eval

# ============ 输出结果 ============
print('========= 问题 4：三架 UAV 各 1 枚的最优遮蔽策略（FY1、FY2、FY3→ M1） =========')
for k, p in enumerate(plan):
    name = f'FY{p.uav + 1}'
    print(f'—— {name} ——')
    print(f'航向角θ = {np.rad2deg(p.theta):.2f}°（相对 +x 逆时针）')
    print(f'飞行速度 v = {p.v:.2f} m/s')
    print(f'投放时刻 t_rel = {p.t_rel:.3f} s')
    print(f'引信延时τ = {p.tau:.3f} s')
    print(f'起爆时刻 t_det = {p.t_det:.3f} s')
    print(f' 投放点 D = ({p.D[0]:.3f}, {p.D[1]:.3f}, {p.D[2]:.3f}) m')
    print(f' 起爆点 C0 = ({p.C0[0]:.3f}, {p.C0[1]:.3f}, {p.C0[2]:.3f}) m\n')

print(f'三枚弹联合遮蔽总时长 = {total_cov:.6f} s')
if total_cov > 60 + 1e-6:
    print('警告: 理论上三弹上限 60 s（各 20 s），当前超过，建议检查网格/步长设置。')